import tkinter as tk
from tkinter import filedialog
import cv2
import torch
from paddleocr import PaddleOCR
import os

# 初始化模型
def init_models():
    # YOLOv5 模型
    yolo_model = torch.hub.load(
        r"C:\Users\23326\Desktop\yolov5-7.0", 
        'custom', 
        path=r"C:\Users\23326\Desktop\yolov5-7.0\runs\train\exp10\weights\best.pt", 
        source='local'
    )
    yolo_model.conf = 0.4  # 设置置信度阈值

    # PaddleOCR 模型
    ocr_model = PaddleOCR(use_angle_cls=True, lang='ch')
    return yolo_model, ocr_model

# 处理图像并检测车牌
def process_image(image_path, yolo_model, ocr_model):
    img = cv2.imread(image_path)
    results = yolo_model(img)
    detections = results.xyxy[0]  # tensor格式: [x1, y1, x2, y2, conf, class]

    # 创建保存目录
    save_dir = r"C:\Users\23326\Desktop\yolov5-7.0\cropped_plates"
    os.makedirs(save_dir, exist_ok=True)

    # 遍历检测结果
    plate_texts = []
    for i, det in enumerate(detections):
        x1, y1, x2, y2 = map(int, det[:4])
        cropped = img[y1:y2, x1:x2]

        # 保存裁剪图像
        save_path = os.path.join(save_dir, f"plate_{i}.jpg")
        cv2.imwrite(save_path, cropped)

        # OCR 识别
        img_rgb = cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB)
        result = ocr_model.ocr(img_rgb, cls=True)
        if result:
            text = " ".join([word_info[1][0] for word_info in result[0]])
            plate_texts.append(text)
        else:
            plate_texts.append("未识别到车牌")

    return plate_texts

# 选择图片并显示结果
def open_image():
    file_path = filedialog.askopenfilename()
    if file_path:
        plate_texts = process_image(file_path, yolo_model, ocr_model)
        result_text = "\n".join(plate_texts) if plate_texts else "未检测到车牌"
        result_label.config(text=result_text)

# 初始化 GUI
root = tk.Tk()
root.title("车牌检测与识别")
root.geometry("600x400")

# 初始化模型
yolo_model, ocr_model = init_models()

# 创建按钮和标签
open_button = tk.Button(root, text="打开图片", command=open_image)
open_button.pack(pady=20)

result_label = tk.Label(root, text="", wraplength=500)
result_label.pack(pady=20)

# 运行 GUI
root.mainloop()